import torch
import matplotlib.pyplot as plt
import numpy as np
import random
from torchvision.utils import make_grid

def visualize_results(model, test_loader, class_names, num_images=10, device=None):
    """
    可视化模型预测结果
    
    参数:
        model (nn.Module): 训练好的模型
        test_loader (DataLoader): 测试数据加载器
        class_names (list): 类别名称列表
        num_images (int): 要可视化的图像数量
        device (torch.device): 设备(CPU/GPU)
    """
    if device is None:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    model.to(device)
    model.eval()
    
    # 获取一批测试数据
    dataiter = iter(test_loader)
    images, labels = next(dataiter)
    images, labels = images.to(device), labels.to(device)
    
    # 随机选择一些图像
    indices = random.sample(range(images.size(0)), num_images)
    images = images[indices]
    labels = labels[indices]
    
    # 获取预测结果
    with torch.no_grad():
        outputs = model(images)
        _, preds = torch.max(outputs, 1)
    
    # 反归一化图像
    mean = torch.tensor([0.4914, 0.4822, 0.4465]).to(device)
    std = torch.tensor([0.2470, 0.2435, 0.2616]).to(device)
    images = images * std.view(1, 3, 1, 1) + mean.view(1, 3, 1, 1)
    
    # 创建图像网格
    img_grid = make_grid(images.cpu(), nrow=5)
    img_grid = img_grid.numpy().transpose((1, 2, 0))
    
    # 创建子图
    plt.figure(figsize=(15, 10))
    plt.imshow(img_grid)
    plt.axis('off')
    
    # 添加预测标签
    for i in range(num_images):
        true_label = class_names[labels[i]]
        pred_label = class_names[preds[i]]
        
        color = 'green' if true_label == pred_label else 'red'
        
        plt.text(
            (i % 5) * 64 + 10, (i // 5) * 64 + 50, 
            f"True: {true_label}\nPred: {pred_label}", 
            color=color, 
            fontsize=12,
            bbox=dict(facecolor='white', alpha=0.7)
        )
    
    plt.tight_layout()
    plt.savefig('predictions.png')
    plt.show()